def test_model(test_loader, model):
    model.eval()
    
    submission = [['ImageId', 'Label']]
    image_id = 1
    for images, _ in test_loader:
        if args.gpu > -1:
            model = model.cuda()
            images = images.cuda()
            log_ps = model(images)
        else:
            model = model.cpu()
            images = images.cpu()
            loglog_psps = model(images)
            
        ps = torch.exp(log_ps)
        top_p, top_class = ps.topk(1, dim=1)
        
        for prediction in top_class:
            submission.append([image_id, prediction.item()])
            image_id += 1
    return submission
submission = test_model(test_loader, model)

submission_df = pd.DataFrame(submission)
submission_df.columns = submission_df.iloc[0]
submission_df = submission_df.drop(0, axis=0)

submission_df.to_csv("submission.csv", index=False)
